#!/usr/bin/env python3
"""
未验证

训练记录：
在二号机上训练
20241121：训练分数达到-39（由于中途误关闭，所以分数应该不准确），测试分数达到0分，继续训练
20241122：由于训练电脑蓝屏，训练分数未知未知，测试分数依旧是昨天的分数，毫无进步，将学习率调低，继续训练一天，无进步则说明该dqn方法无法解决该问题
20241127：训练分数-38分，测试分数0分（产生新的测试分数9分），感觉无任何进步，未调整超参数，是否需要暂缓训练
20250304：完成代码调整，尝试重新训练
在主力机上训练
20250308:Learning Rate: 0.0001，测试分数0分，训练分数-37.770
20250309:Learning Rate: 0.0001，测试分数1.0分，训练分数 -37.170，继续训练
20250310：Learning Rate: 0.0001，测试分数1.0分，训练分数 -36.440，继续训练
20250311:完成代码调整，重新训练，测试分数40分，训练分数 -16.700，看起分数整体是上升趋势，继续训练
20250312：Learning Rate: 0.0001
20250312:测试分数690分，训练分数 48.800，继续训练，看起来进步很大
20250313:Learning Rate: 0.0001，测试分数1286，训练分数277.100，继续训练，整体呈现上升趋势
20250314:Learning Rate: 0.0001，测试分数1286，训练分数366.300，继续训练，整体呈现上升趋势
20250315:Learning Rate: 0.0001，测试分数1286，训练分数384.300，继续训练，整体呈现上升趋势
20250316:Learning Rate: 0.0001，训练到一半中途异常，继续训练
20250317:Learning Rate: 0.0001，测试分数1396，训练分数515.500，继续训练，整体呈现上升趋势
20250318：Learning Rate: 0.0001,测试分数1518，训练分数624.300，继续训练，整体呈现上升趋势
20250319：Learning Rate: 0.0001,测试分数1518，训练分数602.400，继续训练，开放学习率调度器
20250320:Learning Rate: 0.0001，测试分数1518，训练分数625.100，继续训练，关闭学习率调度器，好像有点提升慢
20250321:Learning Rate: 9.025e-05，测试分数1530，训练分数632.9，继续训练，查看学习率调度器
20250322: 暂停训练
20250323: 暂停训练
20250324:Learning Rate: 9.025e-05,测试分数1530.0，训练分数 616.800，继续训练，开放学习率调度器
20250325：Learning Rate: 9.025e-05,测试分数1530，训练分数672.000，继续训练，关闭学习率调度器
20250326：Learning Rate: 7.737809374999998e-05，测试分数1530，训练分数663.300，继续训练，开启学习率调度器
20250327：Learning Rate: 7.737809374999998e-05，开启学习率调度器,测试分数1530，训练分数615.400，继续训练，确认学习率调度器
20250327：Learning Rate: 6.983372960937497e-05，测试分数1530，训练分数650.800，停止训练，play模型，因为没有产生新的测试分数，训练分数也不如之前的高
"""
import gymnasium as gym
import ptan
import argparse
import numpy as np

import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
import os

from tensorboardX import SummaryWriter

from typing import Any
from lib import common
import ale_py

gym.register_envs(ale_py)

class FireResetEnv(gym.Wrapper):
    def __init__(self, env=None):
        """For environments where the user need to press FIRE for the game to start."""
        super(FireResetEnv, self).__init__(env)
        # 以下可知，一些游戏存在FIRE的动作，并且存在FIRE动作的游戏其游戏动作执行有三个以上
        assert env.unwrapped.get_action_meanings()[1] == 'FIRE'
        assert len(env.unwrapped.get_action_meanings()) >= 3

    def step(self, action):
        return self.env.step(action)

    def reset(self, seed: int | None = None, options: dict[str, Any] | None = None):
        # 这里之所以尝试重置后尝试各种动作，是因为不知道哪个是FIRE，继续游戏，所以一个一个尝试
        # 如果不小心游戏结束了，则继续重置
        # 假设游戏继续游戏的按钮在前3
        self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(1)
        if done:
            self.env.reset(seed=seed, options=options)
        obs, _, done, _, info = self.env.step(2)
        if done:
            self.env.reset(seed=seed, options=options)
        return obs, info

class DuelingDQN(nn.Module):
    def __init__(self, input_shape, n_actions):
        super(DuelingDQN, self).__init__()

        self.conv = nn.Sequential(
            nn.Conv2d(input_shape[0], 32, kernel_size=8, stride=4),
            nn.ReLU(),
            nn.Conv2d(32, 64, kernel_size=4, stride=2),
            nn.ReLU(),
            nn.Conv2d(64, 64, kernel_size=3, stride=1),  # 改为64通道，与原论文一致
            nn.ReLU()
        )

        conv_out_size = self._get_conv_out(input_shape)
        
        # 增加一个共享的特征层
        self.fc_common = nn.Sequential(
            nn.Linear(conv_out_size, 512),
            nn.ReLU()
        )
        
        self.fc_adv = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, n_actions)
        )
        
        self.fc_val = nn.Sequential(
            nn.Linear(512, 256),
            nn.ReLU(),
            nn.Linear(256, 1)
        )

    def _get_conv_out(self, shape):
        o = self.conv(torch.zeros(1, *shape))
        return int(np.prod(o.size()))

    def forward(self, x):
        fx = x.float() / 255.0
        conv_out = self.conv(fx).view(fx.size()[0], -1)
        common_features = self.fc_common(conv_out)
        
        val = self.fc_val(common_features)
        adv = self.fc_adv(common_features)
        
        # 正确的Dueling网络组合方式，使用减去平均值
        return val + adv - adv.mean(dim=1, keepdim=True)


class RewardPenaltyWrapper(gym.Wrapper):
    def __init__(self, env, frame_penalty=-0.1, life_loss_penalty=-10):
        super(RewardPenaltyWrapper, self).__init__(env)
        self.frame_penalty = frame_penalty
        self.life_loss_penalty = life_loss_penalty
        self.previous_lives = 0

    def reset(self, **kwargs):
        obs, info = self.env.reset(**kwargs)
        self.previous_lives = info.get('lives', 0)  # 初始生命值
        return obs, info

    def step(self, action):
        obs, reward, done, truncated, info = self.env.step(action)

        # if reward != 0:
        #     reward //= 10.0  # 缩放奖励
        #     if reward <= 0:
        #         reward = 1.0

        # 处理生命减少时的惩罚
        current_lives = info.get('lives', self.previous_lives)
        if current_lives < self.previous_lives:
            reward = self.life_loss_penalty
            self.previous_lives = current_lives
        elif current_lives > self.previous_lives:
            self.previous_lives = current_lives

        return obs, reward, done, truncated, info


def wrap_dqn(env, stack_frames=4, episodic_life=True):
    if episodic_life:
        # 将多条生命的游戏模拟成单条生命ActorCriticAgent
        env = ptan.common.wrappers.EpisodicLifeEnv(env)
    # 增强初始化
    env = ptan.common.wrappers.NoopResetEnv(env, noop_max=30)

    if 'FIRE' in env.unwrapped.get_action_meanings():
        env = FireResetEnv(env)
    env = common.ProcessFrame84(env)
    env = ptan.common.wrappers.ImageToPyTorch(env)
    env = ptan.common.wrappers.FrameStack(env, stack_frames)
    env = RewardPenaltyWrapper(env)
    return env


def test_model(env, net, device, episodes=5):
    with torch.no_grad():
        total_reward = 0.0
        for ep in range(episodes):
            noop_action_count = 0
            pre_action = -1
            obs, _ = env.reset()
            episode_reward = 0.0
            while True:
                obs_v = ptan.agent.default_states_preprocessor([obs]).to(device)
                logits_v = net(obs_v)
                probs = logits_v.data.cpu().numpy()
                action = np.argmax(probs)
                
                # 检测卡住的情况
                if action == 0 and pre_action == action:  # Noop
                    noop_action_count += 1
                    if noop_action_count > 30:
                        break
                else:
                    noop_action_count = 0
                
                pre_action = action
                obs, reward, done, trunc, _ = env.step(action)
                episode_reward += reward
                if done or trunc:
                    break
            
            total_reward += episode_reward
            # print(f"Test episode {ep+1}: reward {episode_reward:.2f}")
            
    return total_reward / episodes


if __name__ == "__main__":
    params = {
        'env_name':         "ALE/BankHeist-v5",
        'stop_reward':      1000.0,
        'run_name':         'BankHeist',
        'replay_size':      200000,  # 增大经验池
        'replay_initial':   20000,   # 增加初始采样量
        'target_net_sync':  500,     # 更频繁地更新目标网络
        'epsilon_frames':   200000,  # 延长探索时间
        'epsilon_start':    1.0,
        'epsilon_final':    0.05,    # 保持一定探索度
        'learning_rate':    0.0001,  # 适中的学习率
        'gamma':            0.99,
        'batch_size':       32       # 减小批量大小，更新更频繁
    }
    parser = argparse.ArgumentParser()
    parser.add_argument("--cuda", default=True, action="store_true", help="Enable cuda")
    parser.add_argument("-n", "--name", required=True, help="Name of the run")
    args = parser.parse_args()
    device = torch.device("cuda" if args.cuda else "cpu")

    save_path = os.path.join("saves", "dueling-" + args.name)
    if not os.path.exists(save_path):
        os.makedirs(save_path)

    env = wrap_dqn(gym.make(params['env_name'], obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)
    test_env = wrap_dqn(gym.make(params['env_name'], obs_type='rgb', frameskip=4, repeat_action_probability=0.0), episodic_life=False)

    writer = SummaryWriter(comment="-" + params['run_name'] + "-dueling")
    # 不同：采用的是竞争DQN网络
    net = DuelingDQN(env.observation_space.shape, env.action_space.n).to(device)
    tgt_net = ptan.agent.TargetNet(net)
    selector = ptan.actions.EpsilonGreedyActionSelector(epsilon=params['epsilon_start'])
    epsilon_tracker = common.EpsilonTracker(selector, epsilon_start=params['epsilon_start'], epsilon_final=params['epsilon_final'], epsilon_frames=params['epsilon_frames'])
    agent = ptan.agent.DQNAgent(net, selector, device=device)

    exp_source = ptan.experience.ExperienceSourceFirstLast(env, agent, gamma=params['gamma'], steps_count=1)
    buffer = ptan.experience.ExperienceReplayBuffer(exp_source, buffer_size=params['replay_size'])
    optimizer = optim.Adam(net.parameters(), lr=params['learning_rate'])
    # 更温和的学习率衰减
    scheduler = optim.lr_scheduler.StepLR(optimizer, step_size=10000, gamma=0.95)

    frame_idx = 0

    if os.path.exists(save_path) and len(os.listdir(save_path)) > 0:
        # 增加加载模型的代码
        checkpoints = sorted(filter(lambda x: "epoch" in x, os.listdir(save_path)),
                             key=lambda x: int(x.split('_')[2].split('.')[0]))
        checkpoint = torch.load(os.path.join(save_path, checkpoints[-1]), map_location=device, weights_only=False)
        frame_idx = checkpoint['frame_idx']
        net.load_state_dict(checkpoint['net'])
        optimizer.load_state_dict(checkpoint['optimizer'])
        scheduler.load_state_dict(checkpoint['scheduler'])
        print("加载模型成功")
        print("")
    # 打印学习率大小
    print("Learning Rate:", scheduler.get_last_lr()[0])

    with common.RewardTracker(writer, params['stop_reward']) as reward_tracker:
        while True:
            frame_idx += 1
            buffer.populate(1)
            epsilon_tracker.frame(frame_idx)

            new_rewards = exp_source.pop_total_rewards()
            if new_rewards:
                if reward_tracker.reward(new_rewards[0], frame_idx, selector.epsilon):
                    break

            if len(buffer) < params['replay_initial']:
                continue

            optimizer.zero_grad()
            batch = buffer.sample(params['batch_size'])
            loss_v = common.calc_loss_dqn(batch, net, tgt_net.target_model, gamma=params['gamma'], device=device)
            loss_v.backward()
            
            # 梯度裁剪，防止爆炸
            torch.nn.utils.clip_grad_norm_(net.parameters(), max_norm=10.0)
            
            optimizer.step()
            # scheduler.step()

            if frame_idx % params['target_net_sync'] == 0:
                for _ in range(12):
                    scheduler.step()

                tgt_net.alpha_sync(0.9)
                checkpoints = {
                    'frame_idx': frame_idx,
                    'net': net.state_dict(),
                    'optimizer': optimizer.state_dict(),
                    'scheduler': scheduler.state_dict(),
                }
                common.save_checkpoints(frame_idx, checkpoints, save_path, "dueling", keep_last=5)

                net.eval()
                test_reward = test_model(test_env, net, device=device, episodes=5)
                net.train()
                print(f"Test reward: {test_reward:.2f}")
                common.save_best_model(test_reward, net.state_dict(), save_path, "a2c-best", keep_best=10)


